#pragma once

#include <torch/torch.h>

class CustomDtaset : public torch::data::Dataset<CustomDtaset>
{
private:
    torch::Tensor inputs_; ///< 输入
    torch::Tensor labels_; ///< 标签

public:
    CustomDtaset(torch::Tensor input_data, torch::Tensor label_data);
    torch::data::Example<> get(size_t index) override;
    torch::optional<size_t> size() const override;
};

class Net : public torch::nn::Module
{
private:
    torch::Tensor matA; //< 状态增益矩阵
    torch::Tensor matB; //< 控制矩阵
    torch::Tensor matH; //< 状态变量对观测变量的增益矩阵

    torch::Tensor KalmanGain; //< 卡尔曼增益

public:
    Net();
    Net(const torch::Tensor& A, const torch::Tensor& B, const torch::Tensor& H, const int m,
        const int n);
    ~Net();

    at::Tensor forward(at::Tensor u, at::Tensor xtm1, at::Tensor z);
};
